Thanks for the PR @young310 — "drop-in Ollama replacement" auto-load is a real gap and the helper-extraction shape is exactly the right architecture. Unfortunately the PR doesn't run as-is. Flagging blockers below.

P0 — blocker (PR is non-functional)

vllm_mlx.server does not export swap_to_model or get_loading_model

vllm_mlx/service/helpers.py:272:

from ..server import get_loading_model, swap_to_model

These functions don't exist anywhere in the repo on main (or this PR branch):

$ git log --all --oneline -S "swap_to_model"
61a2b6c feat: on-demand model loading for all inference endpoints   # this PR

$ grep -rn "def swap_to_model\|def get_loading_model" vllm_mlx/
# (no matches)

Demonstrated runtime failure:

$ python3.12 -c "
import asyncio
from vllm_mlx.service.helpers import ensure_model_loaded
from vllm_mlx.config import get_config
cfg = get_config()
cfg.model_name = 'qwen3.5-4b'
cfg.model_alias = None; cfg.model_path = None; cfg.model_registry = None
asyncio.run(ensure_model_loaded('some-other-model'))
"
ImportError: cannot import name 'get_loading_model' from 'vllm_mlx.server'

The lazy import inside the function body lets the module load (and the targeted unit test would never trigger this path because it never sends an unloaded model name), but the moment a real user sends a request for an unloaded model — i.e., the entire feature this PR claims to add — the server returns 500.

The PR description claims:

Tested locally on macOS (Apple Silicon), rapid-mlx 0.6.30:
# → server swaps model automatically, returns response from Qwen2.5-Coder

That output cannot have come from this codebase. The only "hot-swap" we have today is in the CLI side (vllm_mlx/cli.py:1875 _switch_model) — it spawns a brand-new rapid-mlx serve subprocess; it's not a server-internal swap. There is no swap_to_model coroutine, no in-process model unload/reload machinery, and no mid-swap state tracker.

Two paths forward:

1. Build the missing infra in the same PR. That would require:

   • swap_to_model(model_name) — async, holds a per-process lock, releases the current BatchedEngine (drains in-flight requests, frees Metal allocations), instantiates the new engine, swaps it into cfg.engine. Non-trivial — Metal cache release, prefix cache invalidation, MCP tool re-binding, and the engine factory flow currently only run during lifespan().
   • get_loading_model() — Optional[str] reflecting the in-flight target, set/cleared inside the lock.
   • Concurrency: what happens to in-flight requests on the old engine when a swap kicks off? Drain? Cancel with 503?
   • Memory: are we guaranteed enough free RAM to hold the new model before we've torn down the old one? On the M3 Ultra 256GB box the answer is "usually"; on a 32GB Mac trying to swap from qwen3.5-9b → kimi-48b it's "no".

2. Scope this PR down. Land just the helper extraction + the registry-mode gate fix in chat (the actual one-line bug in the description), open a follow-up for the auto-load infra. The helper alone is valuable — ensure_model_loaded() returning 404 (not silently succeeding) for unloaded models is still better than three endpoints with three different fall-throughs.

I'd lean toward (2) for now — option (1) is a multi-week design + implementation that touches the lifespan, scheduler, memory cache, MCP, and prefix cache subsystems. Worth a design doc + issue first.

P0 — blocker (security / cost)

Auto-loading runs before _validate_model_name

In all three routes:

await ensure_model_loaded(request.model)   # ← this triggers an HF download
_validate_model_name(request.model)        # ← this would have rejected the input

Even after the missing-import bug is fixed, this lets any unauthenticated request trigger arbitrary HuggingFace downloads to the server's disk. A model can be 200GB+. There's no allowlist (cfg.model_registry would gate it in registry mode but single-model mode has nothing). A malicious or misconfigured client can fill the disk in minutes.

Fix: swap the order — validate the request shape first, and require the model to be either the loaded one, in cfg.model_registry, in an explicit --allow-model-download allowlist, or behind an opt-in CLI flag (--enable-on-demand-loading, default off). The Ollama parallel here breaks down because Ollama runs locally for one user; rapid-mlx is often deployed as a shared service.

P1 — should fix

1. HTTP timeout vs download time. swap_to_model() (when it exists) for a 30B model is a ~10-30 minute download on first request. The chat completion request is held open the entire time. Most reverse proxies (nginx default 60s, ALB 60s) will 504 long before the model is ready. Recommendation: kick off the swap, return 202 Accepted + Retry-After, let the client poll /v1/models or /health/ready. Or — if synchronous is required — at least set an explicit cfg.swap_timeout_seconds so users can tune it for their proxy.

2. TOCTOU in the in-flight check.

   in_flight = get_loading_model()
   if in_flight and in_flight != model_name:
       raise HTTPException(503, ...)
   await swap_to_model(model_name)

   Two simultaneous requests for two unloaded models both pass the in_flight is None check, then both call swap_to_model. Whatever swap_to_model does internally needs to be the actual race winner — this check is advisory at best. Recommend documenting that the lock lives in swap_to_model, not here.

3. Anthropic endpoint will try to auto-load claude-* model names. Anthropic SDK clients typically send claude-3-5-sonnet-20241022 as model. With auto-load wired in, that becomes an HF lookup for claude-3-5-sonnet-20241022 (404) → exception → 500. Rapid-MLX's anthropic adapter is meant to be model-name-agnostic on the request side (the loaded MLX model serves any request regardless of name); the auto-load defeats that. Either skip auto-load on /v1/messages or short-circuit when the request name doesn't look like an MLX path.

4. No tests for the new behavior. A change touching three inference endpoints with auto-load semantics needs:

   • _is_model_loaded returns True for current model_name/model_alias/model_path
   • _is_model_loaded returns False for an unrelated string
   • ensure_model_loaded is a no-op when model is loaded
   • ensure_model_loaded raises 503 when a different swap is in flight
   • Each route returns the correct error envelope when swap_to_model itself raises

   The targeted unit test would have caught the missing import in CI before review.

P2 — nits

• _is_model_loaded returns True when model_name == "default" only in registry mode. Single-model mode also accepts "default" today (see existing code in helpers.py:222). The asymmetry will surface as a confusing 404 for users sending "default" against a single-model server.
• Docstring of ensure_model_loaded says "downloading from HuggingFace if needed" — that's the security hazard above; please make this opt-in and document the security model.

Verdict

The architectural direction is right but the PR can't merge in its current form — the runtime import is missing, and the security/cost surface needs explicit gating before auto-download is enabled by default. Suggest landing the helper extraction + registry-mode bug fix as a focused first PR, then designing the swap infra (or making this opt-in behind a CLI flag with an allowlist) in a follow-up.

Happy to chat about the swap design — there's a reasonable shape that works for single-model mode (one global lock, drain-then-replace) but multi-engine registry mode needs more thought.

Reviewed by @raullenchai (Rapid-MLX maintainer).

Code review

Found 1 issue:

1. ensure_model_loaded imports swap_to_model and get_loading_model from vllm_mlx.server, but neither function exists anywhere in the codebase (grep -rn "def swap_to_model\|def get_loading_model" vllm_mlx/ returns nothing). Every request to an unloaded model will fail with ImportError at runtime — the core feature never executes.

https://github.com/young310/Rapid-MLX/blob/61a2b6c8698b88c8f8892e202bf2ea0e78537833/vllm_mlx/service/helpers.py#L252-L264

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

@raullenchai I add some more testing and please have a look, thank you

Review — PR Merge SOP audit

Thanks for the contribution! Did a multi-pass adversarial review against the PR Merge SOP — auto-deploy makes us conservative on external PRs. Below are findings ranked P0 (blocking) / P1 (should fix) / P2 (nit). I've reproduced each against the diff at head=feat/on-demand-model-loading.

P0 — blocking merge

P0-1. Feature flag is unreachable from the standard rapid-mlx serve CLI.
The --enable-on-demand-loading argparse arg lives in vllm_mlx/server.py::main() (legacy entrypoint via python -m vllm_mlx.server). But users invoke rapid-mlx serve <model> which routes through vllm_mlx/cli.py::serve_command() (cli.py:370). That function directly mutates server._api_key / server._default_timeout / etc. (cli.py:425-465) and calls uvicorn.run() itself — it does not call server.main(). The new flag is not wired into the serve subparser (~cli.py:3060) and is never copied onto the singleton from cli.py. So:

• rapid-mlx serve qwen3.5-4b --enable-on-demand-loading → argparse error "unrecognized arguments"
• Even after fixing argparse, serve_command() never sets cfg.enable_on_demand_loading, so ensure_model_loaded() always returns early.

Please wire the flag through cli.py's serve_parser and serve_command. The pattern to mirror is --default-temperature / --default-top-p (already plumbed both ways).

P0-2. swap_to_model does not drain in-flight requests before tearing down the old engine.
server.py:677 acquires _swap_lock and checks _is_model_loaded(model_name) — that protects swap-vs-swap but not swap-vs-request. In single-model mode the sequence is _model_registry.remove(old_name) → await old_engine.stop() (which sets self._loaded=False, drops _model, _engine). Any coroutine currently iterating engine.stream_generate(...) on the old engine will dereference None mid-stream and the client sees a torn SSE response. There's also a window between await old_engine.stop() and load_model(new) where the module global _engine is still pointing at a half-destroyed object — new requests landing there get an opaque AttributeError instead of a clean 503.

Mitigation needs an inflight-request counter (or engine.drain()) before stop(), and _engine = None set explicitly during the window so get_engine() returns 503 cleanly.

P1 — should fix before merge

P1-1. No validation on request.model before it's handed to load_model() → arbitrary HF download / disk-fill DoS.
chat.py:150 and completions.py:45 do await ensure_model_loaded(request.model) before _validate_model_name(request.model). With the flag on, an attacker-controlled string ("evil-attacker/backdoor", "any-public-repo-with-100GB-weights") reaches swap_to_model → load_model(model_name) → huggingface_hub.snapshot_download(...). No HF-repo-name regex, no disk-space precheck, no allowlist. Even with --api-key gating the endpoint, a single authenticated client can fill disk. Need at minimum: regex check ^[\\w\\-./]+/[\\w\\-.]+$, the existing _check_disk_space() call (cli.py:41), and an optional RAPID_MLX_ALLOWED_REPOS env var.

P1-2. _cached_mlx_models() walks the filesystem on every /v1/models request.
routes/models.py:16-43 — no caching, double for walks every snapshot dir. With 50+ cached models this becomes a per-request os.stat() storm an unauthenticated attacker can amplify by polling /v1/models. Add a 30s mtime-keyed lru_cache.

P1-3. Model-type filter is fragile in both directions.
routes/models.py:33 substring set (\"tts\", \"whisper\", \"asr\", \"sentence-transformer\", \"embed\"). Misses: kokoro (TTS), parakeet (ASR), bge-*, e5-*, gte-*, wav2vec, seamless. False positives: Qwen2.5-Embedded-…, anything with tts in the name that's actually a base LM. Detection should look at config.json's architectures field, not substring-match repo names. (We have is_mllm_model() in api/utils.py as a precedent.)

P1-4. The _sync_config double-import workaround leaves module-globals out of sync with the singleton forever.
server.py:725-731 comment is accurate — python -m vllm_mlx.server does create a __main__ instance separate from vllm_mlx.server. But the chosen fix (write to get_config() directly and exclude this field from _sync_config) means _enable_on_demand_loading (module global, used at server.py:1051 for logging) and cfg.enable_on_demand_loading (singleton, used by ensure_model_loaded) can disagree forever. The robust fix is to never define the flag as a module global — store it only on the singleton. Also: cli.py doesn't go through python -m vllm_mlx.server at all, so the workaround is for an entrypoint that isn't the documented one.

P2 — nits

P2-1. No tests added. Per project SOP §3: "every new behavior MUST have a new test". The state machine introduced here has multiple new behaviors that need pinning:

• lock contention (two concurrent swaps coalesce, second acquirer returns fast after _is_model_loaded recheck)
• _loading_model cleared in finally so swap is retryable after load_model failure
• single-model vs registry branch
• _is_model_loaded handles all of {model_name, model_alias, model_path, \"default\"}
• ensure_model_loaded returns early when flag is off (this is the production default path)

P2-2. Apple-Silicon memory hygiene. await old_engine.stop() does not gc.collect() or mx.clear_cache(). On a 64GB box swapping from a 35GB model, the new load can race against unreclaimed Metal allocator buffers. Either add an explicit cleanup, or document why it's unnecessary.

P2-3. Supply-chain audit (per SOP §2.5): clean. No new deps, no workflow changes, no install hooks, no subprocess/urllib.request at import time. Network surface is contained to the existing huggingface_hub.snapshot_download path. The risk is P1-1 (input validation on the download trigger), not new infrastructure.

Summary

The motivation is great — Ollama-style auto-load is exactly the kind of UX win we want. But P0-1 means the feature isn't actually reachable from the supported entrypoint as written, and P0-2 means the swap path is unsafe under concurrent traffic (which is the realistic usage). P1-1 is a meaningful security gap that the flag-off default doesn't help users who legitimately enable the feature.

Marking as request-changes. Happy to discuss the design — particularly for P0-2, whether you'd prefer a drain-counter or a BatchedEngine.drain() API addition. The PR after these fixes will be very close to mergeable.

— Generated with Claude Code (multi-pass adversarial review)

Supplementary findings (pr_validate / DeepSeek V4 Pro pass)

Re-ran a second independent review via pr_validate (DeepSeek V4 Pro). Confirms the core findings above; surfaces 2 edge cases worth fixing alongside the P0s:

P1-5. _cached_mlx_models() is not exception-safe (routes/models.py:23-38).
The iterdir() walks have no try/except. A single broken symlink, permission error, or partially-downloaded snapshot raises OSError / PermissionError and takes down /v1/models entirely (500 instead of returning the entries that did parse). Wrap both the outer iterdir and the inner snapshot walks in try: except OSError: continue.

P1-6. swap_to_model crashes with TypeError: object of type 'NoneType' has no len() if _model_registry is None (server.py:~651).
The branch is registry_mode = len(_model_registry) > 1. _model_registry is currently always initialized to a ModelRegistry instance at module import, but if the swap ever runs before that initialization (during a partially-bootstrapped startup, or after a future refactor) it crashes. Cheap guard: registry_mode = _model_registry is not None and len(_model_registry) > 1.

Both are easy fixes that should land in the same revision as the P0s.

(Tools used: pr_validate pipeline → deepseek-v4-pro review on the full diff. The stress_e2e_bench step also ran but failed for an unrelated infrastructure reason — port 8451 collision on my local box. That's not your PR.)

@raullenchai Thanks for both review rounds and for surfacing #319 — that materially changes how I want to scope this PR.

Force-pushed a scoped-down revision (832d547). Took option (2) from review #1: dropped the swap infra, --enable-on-demand-loading flag, /v1/models cache enumeration, and the __main__ aliasing workaround. What remains is just the helper extraction and the unified 404 path. The swap/lifecycle work belongs in #319 — please treat this PR as a prep step #319 can build on top of, not a competing implementation.

One thing I want to acknowledge first. Review #1 caught that the "Tested locally … server swaps model automatically" line in the v1 description described behavior the code did not implement (swap_to_model did not exist). Review #3 caught the "Tested on Python 3.14" + "460 passed" mismatch against the supported 3.10–3.13 range / 2289+ test count. Both were description errors on my side and you were right to call them out. This revision only claims what I actually ran; I will be stricter about that going forward.

Diff (4 files, ~60 net lines excluding tests):

• vllm_mlx/service/helpers.py: _is_model_loaded() + ensure_model_loaded() (strict 404, no auto-load). Also fixes the P2-1 "default" asymmetry from review feat: Tier 1 optimizations — streaming tool fix, frequency-aware cache, block reuse #2 by routing _validate_model_name() through _is_model_loaded(), so single-mode and registry-mode share one rule.
• vllm_mlx/routes/chat.py + vllm_mlx/routes/completions.py: call ensure_model_loaded() before _validate_model_name().
• vllm_mlx/routes/anthropic.py: unchanged (adapter is model-name-agnostic; SDK clients send claude-* names that shouldn't trigger an HF lookup).
• tests/test_ensure_model_loaded.py: 12 cases — single-mode, registry-mode, "default" in both modes, strict-404 contract.

ensure_model_loaded() is intentionally an extension point: when #319 lands, the lifecycle/swap/drain logic fits inside that helper without changing route code.

Explicitly out of scope (deferred to #319 or a separate PR):

• swap_to_model / get_loading_model / lifecycle (P0-2 review feat: Tier 1 optimizations — streaming tool fix, frequency-aware cache, block reuse #2)
• --enable-on-demand-loading plumbed through cli.py:serve_parser (P0-1)
• in-flight drain + clean _engine window (P0-2)
• await asyncio.to_thread(load_model, ...) to keep the event loop responsive (P0-3 review feat: Prompt cache for SimpleEngine + tool logits safety #3)
• registry-mode swap orphan handling (P0-4)
• 503 TOCTOU under concurrent unknown-model requests (P0-5)
• HF-repo regex / disk-space precheck / allowlist (P1-1)
• /v1/models cache enumeration with caching + arch-based detection + exception safety (P1-2, P1-3, P1-5)
• __main__ aliasing workaround (P1-4)

Verification (only commands I actually ran):

• pytest tests/test_ensure_model_loaded.py → 12 passed
• pytest tests/test_routes.py tests/test_chat_route_tool_tag_leak.py tests/test_anthropic_route_auth.py → 65 passed
• Rebased onto upstream/main at d171d18 (v0.6.50). The earlier conflict against 1cbbb86 is resolved.
• Local env is Python 3.14, which is outside the supported 3.10–3.13 range per pyproject.toml. CI on a supported version is authoritative; I have not made any "full suite passes" claim.

Happy to close this entirely and let the helper live inside #319 if you'd rather not carry two PRs. Also happy to narrow further (e.g. only the P2-1 default-asymmetry fix, dropping the new ensure_model_loaded extension point) — just say the word.

Re-review — scope is right, three small cleanups before merge

Thanks for the tight turn and the explicit deferred-list. The v2 framing — _is_model_loaded extraction + a strict-404 extension point — is exactly the right baseline for #319 to build on. The description honesty + the P2-1 "default" fix are both clean wins. No need to narrow further; please don't close.

You were also right to skip anthropic.py — I almost flagged that as an inconsistency before realizing the Anthropic-compat adapter is model-name-agnostic by design (Claude SDK clients send claude-sonnet-* etc, which would always 404 here). Good call to leave it alone.

Three things to address before this is mergeable:

P1-1. Redundant double-call in chat.py + completions.py

After your patch, routes/chat.py:204-208 runs both:

await ensure_model_loaded(request.model)   # NEW — checks _is_model_loaded
_validate_model_name(request.model)        # OLD — also checks _is_model_loaded

Same in routes/completions.py:45-46. Both call _is_model_loaded() internally and both raise 404 on miss. If the goal is for ensure_model_loaded to be the canonical check (and the extension point #319 builds on), drop _validate_model_name(request.model) from these two routes — keep it imported only for anthropic.py:105 which is unaffected by your change.

P1-2. Error message regression on chat/completions

_validate_model_name's 404 reads "The model \X` does not exist. Available: qwen3.5-4b, phi4-14b, ..."— the "Available" hint is the most useful debugging signal we surface today.ensure_model_loadedraises the shorter"Model `X` is not loaded."` — which, after the fix in P1-1, becomes the only 404 message on chat/completions.

Suggest having ensure_model_loaded mirror the longer format:

async def ensure_model_loaded(model_name: str | None) -> None:
    if _is_model_loaded(model_name):
        return
    cfg = get_config()
    available = (
        ", ".join(cfg.model_registry.list_model_names())
        if cfg.model_registry is not None
        else cfg.model_name or "<none>"
    )
    raise HTTPException(
        status_code=404,
        detail=f"The model `{model_name}` does not exist. Available: {available}",
    )

Then _validate_model_name and ensure_model_loaded become genuinely interchangeable on the strict-404 path, and #319's swap logic can replace the raise cleanly.

P1-3. Edge-case behavior change: server with no model configured

_validate_model_name had an early return for not cfg.model_name and cfg.model_registry is None (returning silently when nothing is configured). _is_model_loaded in your refactor returns False for that same case, so ensure_model_loaded raises 404 — a slight behavior change for the unconfigured-server case.

In practice the server always has a model_name post-CLI parse, so this is theoretical, but if you adopt P1-2's longer-message form, the if not cfg.model_name and cfg.model_registry is None short-circuit naturally lives at the top of ensure_model_loaded too — keeps behavior parity. Up to you whether it's worth covering.

Nits

• tests/test_ensure_model_loaded.py:42 — MagicMock with __contains__ = lambda self, x: ... works but Mock(spec=ModelRegistry) would catch interface drift if the registry contract grows. Optional.
• The 65-test local pass on Python 3.14 is enough signal that nothing obvious broke; CI on 3.10/3.11/3.12 will run as soon as I approve the workflow trigger for your fork (external-contributor gating). Will do that on the next push.

Concrete next step

Push a patch that:

1. Drops _validate_model_name(request.model) from routes/chat.py:208 and routes/completions.py:46
2. Updates ensure_model_loaded to emit the "Available: ..." 404 detail
3. Optional: adds the unconfigured-server short-circuit for P1-3
4. Optional: a test case asserting the 404 detail contains Available: to lock in the contract

When that lands, I'll trigger CI and we should be green to merge.

Re: your offer to fold the helper into #319 — keep this PR. #319 is much bigger and the helper has standalone value (the P2-1 default fix alone is worth merging). Treating this as #319's prerequisite is the right framing.

Thanks again for the careful iteration.